from typing import Any, TypeVar

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, resume

AnySample = TypeVar('AnySample')


class CountedDataset(Dataset[AnySample]):
    """ A Dataset which is estimated to have the given number of samples

    Please notice that the dataset can actually have different number of samples.
    The estimated number only affects `__len__()` for progress indicator like `tqdm`.
    To actually limit the number of samples, use `BoundedDataset` instead.
    """

    def __init__(self, dataset: Dataset[AnySample], samples_n: int):
        self._dataset = dataset
        self._samples_n = samples_n

    @property
    def dataset(self) -> Dataset[AnySample]:
        return self._dataset

    def __len__(self) -> int:
        return self._samples_n

    def __getitem__(self, index: int, /) -> AnySample:
        return self._dataset[index]

    def __iter__(self) -> DatasetIterator[AnySample]:
        return iter(self._dataset)

    def __resume__(self, state: Any | None = None, /) -> DatasetIterator[AnySample]:
        return resume(self._dataset, state)
